from collections.abc import Iterable, Iterator
from typing import Generic, TypeVar

from zkl_aiutils_datasets.wrapping import wrap_dataset, wrap_dataset_iterator
from .counter import CounterDataset, CounterIterator
from .zipped import ZippedDataset, ZippedIterator

AnySample = TypeVar("AnySample")


class IndexedIterator(ZippedIterator[AnySample, int], Generic[AnySample]):
    def __init__(self, iterator: Iterator[AnySample] | Iterable[AnySample], *, start: int = 0):
        super().__init__((wrap_dataset_iterator(iterator), CounterIterator(start)))

    def __next__(self) -> tuple[AnySample, int]:
        return super().__next__()

    def __skip__(self, samples_n: int):
        super().__skip__(samples_n)


class IndexedDataset(ZippedDataset[AnySample, int], Generic[AnySample]):
    def __init__(self, dataset: Iterable[AnySample], *, start: int = 0):
        super().__init__((wrap_dataset(dataset), CounterDataset(start)))

    def __len__(self) -> int:
        return len(self._datasets[0])

    def __getitem__(self, index: int, /) -> tuple[AnySample, int]:
        return super().__getitem__(index)

    def __iter__(self) -> ZippedIterator[AnySample, int]:
        return super().__iter__()

    def __resume__(self, state: tuple | None = None, /) -> ZippedIterator[AnySample, int]:
        return super().__resume__(state)
